#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/err.h>

void handle_errors() {
    ERR_print_errors_fp(stderr);
    abort();
}

void sign_file(const char *private_key_file, const char *input_file, const char *output_file) {
    FILE *fp = fopen(input_file, "rb");
    if (!fp) {
        perror("Unable to open input file");
        exit(EXIT_FAILURE);
    }

    // 读取文件内容
    fseek(fp, 0, SEEK_END);
    long file_size = ftell(fp);
    fseek(fp, 0, SEEK_SET);
    unsigned char *data = malloc(file_size);
    if (!data) {
        perror("Unable to allocate memory");
        fclose(fp);
        exit(EXIT_FAILURE);
    }
    fread(data, 1, file_size, fp);
    fclose(fp);

    // 读取私钥
    FILE *key_fp = fopen(private_key_file, "r");
    if (!key_fp) {
        perror("Unable to open private key file");
        free(data);
        exit(EXIT_FAILURE);
    }

    // 从 PEM 文件读取 SM2 私钥
    EVP_PKEY *pkey = PEM_read_PrivateKey(key_fp, NULL, NULL, NULL);
    fclose(key_fp);
    if (!pkey) {
        handle_errors();
    }

    // 创建签名上下文
    EVP_MD_CTX *ctx = EVP_MD_CTX_new();
    if (!ctx) {
        handle_errors();
    }

    // 初始化签名操作
    if (EVP_DigestSignInit(ctx, NULL, EVP_sm3(), NULL, pkey) != 1) {
        handle_errors();
    }

    // 提供要签名的数据
    if (EVP_DigestSignUpdate(ctx, data, file_size) != 1) {
        handle_errors();
    }

    // 获取签名所需的缓冲区大小
    size_t sig_len;
    if (EVP_DigestSignFinal(ctx, NULL, &sig_len) != 1) {
        handle_errors();
    }

    // 分配内存用于签名
    unsigned char *sig = malloc(sig_len);
    if (!sig) {
        perror("Unable to allocate memory for signature");
        EVP_MD_CTX_free(ctx);
        EVP_PKEY_free(pkey);
        free(data);
        exit(EXIT_FAILURE);
    }

    // 获取签名
    if (EVP_DigestSignFinal(ctx, sig, &sig_len) != 1) {
        handle_errors();
    }

    // 将签名写入输出文件
    FILE *out_fp = fopen(output_file, "wb");
    if (!out_fp) {
        perror("Unable to open output file");
        free(sig);
        EVP_MD_CTX_free(ctx);
        EVP_PKEY_free(pkey);
        free(data);
        exit(EXIT_FAILURE);
    }

    fwrite(sig, 1, sig_len, out_fp);
    fclose(out_fp);

    printf("Signature successfully written to %s\n", output_file);

    // 清理资源
    free(sig);
    EVP_MD_CTX_free(ctx);
    EVP_PKEY_free(pkey);
    free(data);
}

int main(int argc, char **argv) {
    if (argc != 4) {
        fprintf(stderr, "Usage: %s <private_key.pem> <input_file> <output_file>\n", argv[0]);
        return EXIT_FAILURE;
    }

    sign_file(argv[1], argv[2], argv[3]);
    return EXIT_SUCCESS;
}
